import os
from dotenv import load_dotenv
from typing import Annotated, List, Literal, Union

from langchain_core.messages import ToolMessage
from langchain_core.tools import InjectedToolCallId, tool
from langgraph.prebuilt import InjectedState
from langgraph.prebuilt.chat_agent_executor import AgentState
from langgraph.types import Command
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent
from utils import format_messages

import json
from langchain_core.messages import messages_to_dict

@tool
def calculator(
        operation: Literal["add", "subtract", "multiply", "divide"],
        a: Union[int, float],
        b: Union[int, float],
) -> Union[int, float]:
    """定义一个 calculator tool.

    Arg:
        operation (str): The operation to perform ('add', 'subtract', 'multiply', 'divide').
        a (float or int): The first number.
        b (float or int): The second number.

    Returns:
        result (float or int): 运算结果
    Example
        Divide: result   = a / b
        Subtract: result = a - b
    """
    if operation == 'divide' and b == 0:
        return {"error": "分母不能为0"}

    # Perform calculation
    if operation == 'add':
        result = a + b
    elif operation == 'subtract':
        result = a - b
    elif operation == 'multiply':
        result = a * b
    elif operation == 'divide':
        result = a / b
    else:
        result = "unknown operation"
    return result


def reduce_list(left: list | None, right: list | None) -> list:
    """安全地合并两个列表，能够处理任一或两个输入为 None 的情况。

    args:
        left (list | None): 要合并的第一个列表，或 None。
        right (list | None): 要合并的第二个列表，或 None。

    returns:
        list: 一个包含两个输入列表所有元素的新列表。
              如果某个输入为 None，则将其视为空列表处理。
    """
    if not left:
        left = []
    if not right:
        right = []
    return left + right

class CalcState(AgentState):
    """graph state"""
    ops: Annotated[List[str], reduce_list]


@tool
def calculator_wstate(
        operation: Literal["add", "subtract", "multiply", "divide"],
        a: Union[int, float],
        b: Union[int, float],
        state: Annotated[CalcState, InjectedState],  # not sent to LLM
        tool_call_id: Annotated[str, InjectedToolCallId]  # not sent to LLM
) -> Union[int, float]:
    """Define a two-input calculator tool.

    Arg:
        operation (str): The operation to perform ('add', 'subtract', 'multiply', 'divide').
        a (float or int): The first number.
        b (float or int): The second number.

    Returns:
        result (float or int): the result of the operation
    Example
        Divide: result   = a / b
        Subtract: result = a - b
    """
    if operation == 'divide' and b == 0:
        return {"error": "Division by zero is not allowed."}

    # Perform calculation
    if operation == 'add':
        result = a + b
    elif operation == 'subtract':
        result = a - b
    elif operation == 'multiply':
        result = a * b
    elif operation == 'divide':
        result = a / b
    else:
        result = "unknown operation"
    ops = [f"({operation}, {a}, {b}),"]
    return Command(
        update={
            "ops": ops,
            "messages": [
                ToolMessage(f"{result}", tool_call_id=tool_call_id)
            ]
        }
    )

def main():
    load_dotenv(os.path.join("..", ".env"), override=True)
    model = os.getenv("MODEL")
    base_url = os.getenv("BASE_URL")
    api_key = os.getenv("_API_KEY")
    # Create agent using create_react_agent directly
    SYSTEM_PROMPT = "你是一位乐于助人的算术助手，擅长使用计算器。"
    model = ChatOpenAI(model=model, base_url=base_url, api_key=api_key)
    tools = [calculator]

    # create agent
    agent = create_react_agent(model, tools, prompt=SYSTEM_PROMPT,
                               # state_schema=AgentState,  # 默认
                               ).with_config({"recursion_limit": 20})  # recursion_limit 限制了agent运行的步骤数。

    # # show the agent
    # png_data = agent.get_graph(xray=True).draw_mermaid_png()
    # output_path = "agent_graph.png"  # 保存到文件
    # with open(output_path, "wb") as f:
    #     f.write(png_data)

    # example usage
    result1 = agent.invoke(
        {
            "messages": [
                {
                    "role": "user",
                    "content": "3.1 * 4.2 的结果等于多少？"
                }
            ]
        }
    )
    format_messages(result1["messages"])

    messages_dict = messages_to_dict(result1["messages"])
    print(json.dumps({"messages": messages_dict}, indent=2, ensure_ascii=False))

def custom_state():
    import json
    from langchain_core.messages import messages_to_dict
    load_dotenv(os.path.join("..", ".env"), override=True)
    model = os.getenv("MODEL")
    base_url = os.getenv("BASE_URL")
    api_key = os.getenv("_API_KEY")

    SYSTEM_PROMPT = "你是一位乐于助人的算术助手，擅长使用计算器。"
    model = ChatOpenAI(model=model, base_url=base_url, api_key=api_key)
    tools = [calculator_wstate]  # new tool

    # Create agent
    agent = create_react_agent(
        model,
        tools,
        prompt=SYSTEM_PROMPT,
        state_schema=CalcState,  # now defining state scheme
    ).with_config({"recursion_limit": 20})

    # Example usage
    result2 = agent.invoke(
        {
            "messages": [
                {
                    "role": "user",
                    "content": "3.1 * 4.2 的结果等于多少？",
                }
            ],
        }
    )
    # format_messages(result2["messages"])

    # 注意，现在响应中包含了 `ops` 字段
    # messages_dict = messages_to_dict(result2["messages"])
    # ops_value = result2["ops"]
    # print(json.dumps({"messages": messages_dict}, indent=2, ensure_ascii=False))
    # print(json.dumps({"ops": ops_value}, indent=2, ensure_ascii=False))

    a = "让我们再尝试一个例子。请注意，本例中包含了**双重工具调用**（dual tool call），工具节点将会**并行执行**这两个调用。"
    # Example usage
    result3 = agent.invoke(
        {
            "messages": [
                {
                    "role": "user",
                    "content": "3.1 * 4.2 + 5.5 * 6.5的结果等于多少？",
                }
            ],
        }
    )
    format_messages(result3["messages"])
    messages_dict3 = messages_to_dict(result3["messages"])
    ops_value3 = result3["ops"]
    print(json.dumps({"messages": messages_dict3}, indent=2, ensure_ascii=False))
    print(json.dumps({"ops": ops_value3}, indent=2, ensure_ascii=False))

if __name__ == '__main__':
    custom_state()